package mcpclient

import (
	"context"
	"log/slog"
	"sync"
	"sync/atomic"

	"github.com/mark3labs/mcp-go/client"
	"github.com/mark3labs/mcp-go/client/transport"
	"github.com/mark3labs/mcp-go/mcp"
)

type SseClient struct {
	name      string
	config    SseServerConfig
	client    *client.Client
	mu        sync.Mutex
	taskCount atomic.Int32
	tools     []mcp.Tool
	open      bool
}

func NewSseClient(name string, config SseServerConfig) *SseClient {
	slog.Debug("NewSseClient", "name", name)
	sse, err2 := transport.NewSSE(config.Url)
	if err2 != nil {
		slog.Error("NewSse", "err", err2)
	}

	cli := client.NewClient(sse)

	s := &SseClient{
		name:      name,
		config:    config,
		client:    cli,
		taskCount: atomic.Int32{},
	}

	err := s.connect()
	if err != nil {
		slog.Error("SseClient connect", "err", err)
	}

	err = s.initTools()
	if err != nil {
		slog.Error("NewSseClient", "init tools error", err)
	}

	return s
}

func (s *SseClient) connect() error {
	s.mu.Lock()
	defer s.mu.Unlock()

	if s.open {
		return nil
	}
	slog.Debug("SseClient.connect", "name", s.name)

	ctx := context.Background()
	err := s.client.Start(ctx)
	if err != nil {
		slog.Error("Error start server", "error", err)
		return err
	}

	request := mcp.InitializeRequest{}
	initResult, err := s.client.Initialize(ctx, request)
	if err != nil {
		slog.Error("Error initializing server", "error", err)
		return err
	}
	slog.Info("Successfully initialized server", "initResult", initResult)
	s.open = true
	return nil
}

func (s *SseClient) disconnect() error {
	s.mu.Lock()
	defer s.mu.Unlock()

	if !s.open {
		return nil
	}
	slog.Debug("SseClient.disconnect", "name", s.name)
	if s.taskCount.Load() > 0 {
		slog.Debug("sseClient having task, not close", "taskCount", s.taskCount.Load())
		return nil
	}

	err := s.client.Close()
	if err != nil {
		slog.Error("Error closing server", "error", err)
		return err
	}
	slog.Debug("disconnect finished", "name", s.name)
	s.open = false
	return nil
}

func (s *SseClient) initTools() error {
	slog.Debug("SseClient.initTools", "name", s.name)

	request := mcp.ListToolsRequest{}
	toolResult, err := s.client.ListTools(context.Background(), request)
	if err != nil {
		slog.Error("Error listing tools", "error", err)
		return err
	}
	slog.Debug("Successfully listed tools", "toolResult", toolResult)
	s.tools = toolResult.Tools
	return nil
}

func (s *SseClient) CallTool(request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	slog.Debug("SseClient.CallTool", "name", s.name, "request", request)

	toolResult, err := s.client.CallTool(context.Background(), request)
	if err != nil {
		slog.Error("Error calling tool", "error", err)
		return nil, err
	}
	slog.Debug("Successfully calling tool", "toolResult", toolResult)
	return toolResult, nil
}
